import random
import unittest

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

from training.models.loss import FocalLossV2

torch.manual_seed(15)
random.seed(15)
np.random.seed(15)
torch.backends.cudnn.deterministic = True


class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()
        net = torchvision.models.resnet18(pretrained=False)
        self.conv1 = net.conv1
        self.bn1 = net.bn1
        self.maxpool = net.maxpool
        self.relu = net.relu
        self.layer1 = net.layer1
        self.layer2 = net.layer2
        self.layer3 = net.layer3
        self.layer4 = net.layer4
        self.out = nn.Conv2d(512, 3, 3, 1, 1)

    def forward(self, x):
        feat = self.conv1(x)
        feat = self.bn1(feat)
        feat = self.relu(feat)
        feat = self.maxpool(feat)
        feat = self.layer1(feat)
        feat = self.layer2(feat)
        feat = self.layer3(feat)
        feat = self.layer4(feat)
        feat = self.out(feat)
        out = F.interpolate(feat, x.size()[2:], mode='bilinear', align_corners=True)
        return out


class MyFocalLossTestCase(unittest.TestCase):

    def test_focal_loss_v1(self):
        net1 = Model()
        net2 = Model()
        net2.load_state_dict(net1.state_dict())

        criteria1 = FocalLossV2()
        net1.cuda()
        net2.cuda()
        net1.train()
        net2.train()
        criteria1.cuda()

        optim1 = torch.optim.SGD(net1.parameters(), lr=1e-2)
        optim2 = torch.optim.SGD(net2.parameters(), lr=1e-2)

        bs = 2
        for it in range(300000):
            inten = torch.randn(bs, 3, 224, 244).cuda()
            lbs = torch.randint(0, 2, (bs, 3, 224, 244)).cuda()
            logits = net1(inten)
            loss1 = criteria1(logits, lbs)
            optim1.zero_grad()
            loss1.backward()
            optim1.step()
            logits = net2(inten)
            optim2.zero_grad()
            optim2.step()
            with torch.no_grad():
                if (it + 1) % 50 == 0:
                    print('iter: {}, ================='.format(it + 1))
                    print('out.weight: ', torch.mean(torch.abs(net1.out.weight - net2.out.weight)).item())
                    print('conv1.weight: ', torch.mean(torch.abs(net1.conv1.weight - net2.conv1.weight)).item())
                    print('loss: ', loss1.item() - loss2.item())


if __name__ == '__main__':
    unittest.main()
